
import os
import cv2
import shutil


frames_dirs = [
    'C:/Users/jario/Downloads/driver_161_90frame/06030819_0755.MP4',
    'C:/Users/jario/Downloads/driver_161_90frame/06030831_0759.MP4',
]
gray_dir = 'C:/dataset/A2RL/ws_lane_240403/Yas_Grey_240322'


def show_lane_in_dir(frames_dir):
    file_names = os.listdir(frames_dir)
    father_dir = os.path.dirname(frames_dir)

    colors = [(0, 0, 255), (0, 255, 255), (255, 0, 0), (0, 255, 0)]

    for file_name in file_names:
        if os.path.splitext(file_name)[-1] == '.jpg':
            img_fn = os.path.join(frames_dir, file_name)
            ann_fn = os.path.join(frames_dir, os.path.splitext(file_name)[0] + '.lines.txt')
            if not os.path.exists(ann_fn):
                os.remove(img_fn)
                continue
            if not os.path.exists(img_fn):
                os.remove(ann_fn)
                continue
            img = cv2.imread(img_fn)
            # print(img_fn)
            img2 = img.copy()
            lanes = []
            if not os.path.exists(ann_fn):
                os.remove(img_fn)
                continue
            with open(ann_fn) as f:
                lines = f.readlines()
                for line in lines:
                    if len(line) > 0:
                        print(line)
                        line = [int(round(float(l))) for l in line.strip().split(' ') if len(l) > 0]
                        lane = []
                        for i in range(0, len(line), 2):
                            lane.append([line[i], line[i + 1]])
                        lanes.append(lane)

            for i, lane in enumerate(lanes):
                for pt in lane:
                    cv2.circle(img, pt, 3, colors[i], -1, cv2.LINE_AA)
                    cv2.circle(img, pt, 3, (0, 0, 0), 1, cv2.LINE_AA)
            # print(img_fn)
            # print(ann_fn)
            """
            img = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
            file_name = os.path.splitext(file_name)[0]
            opath = os.path.join(gray_dir, file_name + '_Gray.jpg')
            cv2.imwrite(opath, img)
            shutil.copy(ann_fn, os.path.join(gray_dir, file_name + '_Gray.lines.txt'))
            """

            cv2.imshow('img', img)
            cv2.waitKey(100)


def write_txt(fp):
    sub_dirs = os.listdir(fp)
    with open(os.path.join(fp, 'list', 'train.txt'), 'w') as f:
        for sub_dir in sub_dirs:
            img_names = os.listdir(os.path.join(fp, sub_dir))
            for img_name in img_names:
                if os.path.splitext(img_name)[1] == '.jpg':
                    to_out = '/' + sub_dir + '/' + img_name
                    f.write(to_out + '\r')


def to_vid(path):
    img_fns = os.listdir(path)
    for img_fn in img_fns:
        if os.path.splitext(img_fn)[1] == '.jpg':
            img = cv2.imread(os.path.join(path, img_fn))
            print(img_fn)


if __name__ == '__main__':
    # to_vid('C:/dataset/A2RL/ws_lane_wg_240406/Yas_Grey_240322')
    show_lane_in_dir(r'C:\dataset\A2RL\yas_lane_tusimple_240414_2\images_from_video_v240415_2')
    # write_txt('C:/dataset/A2RL/ws_lane_wg_240403')
